import re
import subprocess
from HostManager import SSH
from AppManager import Crontab

FireWallDict = {"Op主机": {"data": ""}}


# 获取OP所在本机的防火墙
def getSelfFireWallDict():
    global FireWallDict
    tmp = "firewall-cmd --list-ports"
    p = subprocess.Popen(tmp, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    res = p.stdout.read().decode("utf-8")
    if res == "FirewallD is not running\n":
        res = ""
    FireWallDict["Op主机"]["data"] = str2fwdict(res)
    return


# 修改OP所在本机的防火墙
def EditSelfFW(Rules, Type):
    if Type == "add":
        tmp = "firewall-cmd --zone=public --add-port=%s/%s %s" % (Rules["port"], Rules["protocol"], Rules["params"])
    else:
        tmp = "firewall-cmd --zone=public --remove-port=%s/%s %s" % (Rules["port"], Rules["protocol"], Rules["params"])
    p = subprocess.Popen(tmp, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    res = p.stdout.read().decode("utf-8")
    if res == "FirewallD is not running\n":
        res = ""
    subprocess.Popen("firewall-cmd --reload", shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    return res


# 获取已有资产的防火墙
def getAllFireWallDict():
    global FireWallDict
    tmp = "firewall-cmd --list-ports"
    for hostname, hostdata in Crontab.hostdict.items():
        FireWallDict[hostname] = {}
        command = [tmp]
        try:
            res = SSH.SshComm(hostdata["ip"], hostdata["port"], hostdata["user"], hostdata["password"], command)
            FireWallDict[hostname]["data"] = str2fwdict(res[tmp]["stdout"])
        except Exception as e:
            print("Firewall:无法连接", type(e), e)
            FireWallDict[hostname]["data"] = {}


# 修改已有资产的防火墙
def EditFW(hostname, Rules, Type):
    if Type == "add":
        tmp = "firewall-cmd --zone=public --add-port=%s/%s %s" % (Rules["port"], Rules["protocol"], Rules["params"])
    else:
        tmp = "firewall-cmd --zone=public --remove-port=%s/%s %s" % (Rules["port"], Rules["protocol"], Rules["params"])
    command = [tmp, "firewall-cmd --reload"]
    hostdata = Crontab.hostdict[hostname]
    res = SSH.SshComm(hostdata["ip"], hostdata["port"], hostdata["user"], hostdata["password"], command)
    return res[tmp]["stdout"]


# 将乱糟糟的返回格式化
def str2fwdict(string):
    if string:
        tmpdict = {}
        i = 1
        string = string.replace("\n", "")
        res = re.split("[ ]", string)
        for tmp in res:
            tmpres = re.split("/", tmp)
            tmpdict[str(i)] = {
                "type": "ACCEPT",
                "port": tmpres[0],
                "protocol": tmpres[1]
            }
            i += 1
        return tmpdict
    else:
        return ""


# 修改防火墙规则
def EditFWRule(hostname, Rules, Type):
    if not Type == "add":
        ID = Rules
        Rules = {
            "port": FireWallDict[hostname]["data"][ID]["port"],
            "protocol": FireWallDict[hostname]["data"][ID]["protocol"],
            "type": FireWallDict[hostname]["data"][ID]["type"],
            "params": "--permanent"
        }
    if hostname == "Op主机":
        res = EditSelfFW(Rules, Type)
    else:
        res = EditFW(hostname, Rules, Type)
    getSelfFireWallDict()
    getAllFireWallDict()
    return res


getSelfFireWallDict()
getAllFireWallDict()
